(*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *)

open! NS0
include NSMap_intf

type ('key, +'a, 'compare_key) t = ('key, 'a, 'compare_key) OMap.t [@@deriving compare, equal, sexp]

type ('compare_key, 'compare_a) compare = ('compare_key, 'compare_a) OMap.compare
[@@deriving compare, equal, sexp]

module Make_from_Comparer (Key : sig
  type t [@@deriving equal, sexp_of]

  include NSComparer.S with type t := t
end) =
struct
  module M = OMap.Make (Key)

  type key = Key.t

  type compare_key = Key.compare

  type 'a t = 'a M.t [@@deriving compare]

  type 'compare_a compare = 'compare_a M.compare [@@deriving compare, equal, sexp]

  let comparer = M.comparer

  include M.Provide_equal (Key)
  include M.Provide_sexp_of (Key)
  module Provide_of_sexp = M.Provide_of_sexp
  include Iter.Map.Adapt (M)

  let empty = M.empty

  let singleton = M.singleton

  let add_exn ~key ~data m =
    assert (not (M.mem key m)) ;
    M.add key data m


  let add ~key ~data m = M.add key data m

  let add_absent ~key ~data m =
    M.update key (function Some data -> Some data | None -> Some data) m


  let add_multi ~key ~data m =
    M.update key (function Some vs -> Some (data :: vs) | None -> Some [data]) m


  let remove key m = M.remove key m

  let merge l r ~f =
    let combine k lo ro =
      match (lo, ro) with
      | Some lv, Some rv ->
          f k (`Both (lv, rv))
      | Some lv, None ->
          f k (`Left lv)
      | None, Some rv ->
          f k (`Right rv)
      | None, None ->
          None
    in
    M.merge combine l r


  let merge_endo l r ~f =
    let change = ref false in
    let l' =
      merge l r ~f:(fun key side ->
          let f_side = f key side in
          ( match (side, f_side) with
          | (`Both (data, _) | `Left data), Some data' when data' == data ->
              ()
          | _ ->
              change := true ) ;
          f_side )
    in
    if !change then l' else l


  let union x y ~f = M.union f x y

  let union_absent t u =
    let change = ref false in
    let t' =
      M.merge
        (fun _ v1 v2 ->
          match v1 with
          | Some _ ->
              v1
          | None ->
              change := true ;
              v2 )
        t u
    in
    if !change then t' else t


  let partition m ~f = M.partition f m

  let partition_map m ~f =
    M.fold
      (fun k v (l, r) ->
        match (f k v : _ Either.t) with Left a -> (M.add k a l, r) | Right b -> (l, M.add k b r) )
      m (empty, empty)


  let is_empty = M.is_empty

  let is_singleton = M.is_singleton

  let length = M.cardinal

  let only_binding = M.only_binding

  let classify = M.classify

  let choose = M.choose_opt

  let choose_exn = M.choose

  let min_binding = M.min_binding_opt

  let max_binding = M.max_binding_opt

  let max_binding_exn = M.max_binding

  let mem k m = M.mem k m

  let find_exn k m = M.find k m

  let find k m = M.find_opt k m

  let find_multi k m = match M.find_opt k m with None -> [] | Some vs -> vs

  let find_update x m ~f =
    let found = ref None in
    let m =
      M.update x
        (function
          | None ->
              f None
          | some_v ->
              found := some_v ;
              f some_v )
        m
    in
    (!found, m)


  let find_and_remove = find_update ~f:(fun _ -> None)

  let find_or_add k v m = find_update k ~f:(function None -> Some v | some_v -> some_v) m

  let find_or_add_lazy k m ~f =
    let result = ref None in
    let m =
      M.update k
        (function
          | None ->
              let v' = f () in
              result := Some (`Added v') ;
              Some v'
          | Some v ->
              result := Some (`Found v) ;
              Some v )
        m
    in
    (NSOption.get_exn !result, m)


  let pop_min_binding m = min_binding m |> NSOption.map ~f:(fun (k, v) -> (k, v, remove k m))

  let update k m ~f = M.update k f m

  let map m ~f = M.map f m

  let mapi m ~f = M.mapi (fun key data -> f ~key ~data) m

  let map_endo t ~f = map_endo map t ~f

  let filter_mapi m ~f = M.filter_map (fun key data -> f ~key ~data) m

  let iter m ~f = M.iter (fun _ data -> f data) m

  let iteri m ~f = M.iter (fun key data -> f ~key ~data) m

  let existsi m ~f = M.exists (fun key data -> f ~key ~data) m

  let for_alli m ~f = M.for_all (fun key data -> f ~key ~data) m

  let fold m s ~f = M.fold (fun key data acc -> f ~key ~data acc) m s

  let fold_until (type res) m s ~f ~finish =
    let state = ref s in
    let exception Stop of res in
    try
      iteri m ~f:(fun ~key ~data ->
          match f ~key ~data !state with
          | `Continue s ->
              state := s
          | `Stop r ->
              raise_notrace (Stop r) ) ;
      finish !state
    with Stop r -> r


  let symmetric_diff l r ~eq =
    let seq = ref Iter.empty in
    let yield x = seq := Iter.cons x !seq in
    merge l r ~f:(fun k vv ->
        ( match vv with
        | `Both (lv, rv) when eq lv rv ->
            ()
        | `Both vv ->
            yield (k, `Unequal vv)
        | `Left lv ->
            yield (k, `Left lv)
        | `Right rv ->
            yield (k, `Right rv) ) ;
        None )
    |> ignore ;
    !seq


  let pp pp_k pp_v fs m =
    Format.fprintf fs "@[<1>[%a]@]"
      (NSList.pp ",@ " (fun fs (k, v) -> Format.fprintf fs "@[%a@ @<2>↦ %a@]" pp_k k pp_v v))
      (Iter.to_list (to_iter m))


  let pp_diff ?(pre = ("[@[<hv>" : (unit, unit) fmt)) ?(suf = ("@]];@ " : (unit, unit) fmt))
      ?(sep = (";@ " : (unit, unit) fmt)) pp_key pp_val pp_diff_val ~eq fs (x, y) =
    let pp_diff_elt fs = function
      | k, `Left v ->
          Format.fprintf fs "-- [@[%a@ @<2>↦ %a@]]" pp_key k pp_val v
      | k, `Right v ->
          Format.fprintf fs "++ [@[%a@ @<2>↦ %a@]]" pp_key k pp_val v
      | k, `Unequal vv ->
          Format.fprintf fs "[@[%a@ @<2>↦ %a@]]" pp_key k pp_diff_val vv
    in
    let sd = Iter.to_list (symmetric_diff ~eq x y) in
    NSList.pp ~pre ~suf sep pp_diff_elt fs sd
end
[@@inline]

module Make (Key : sig
  type t [@@deriving compare, equal, sexp_of]
end) =
Make_from_Comparer (struct
  include Key
  include NSComparer.Make (Key)
end)
[@@inline]
